//! Type inference for NRC language

use crate::ast::*;
use crate::error::Result;
use std::collections::HashMap;

use super::TypeEnvironment;

mod declarations;
mod expressions;
mod statements;

/// Type inference engine for NRC language
pub struct TypeInference {
    /// Type environment
    environment: TypeEnvironment,
    /// Type constraints
    constraints: Vec<TypeConstraint>,
    /// Inferred types
    inferred_types: HashMap<String, Type>,
    /// Stack of try contexts for propagating `?`/`!`
    try_context_stack: Vec<Option<TryContext>>,
}

/// Type constraint
#[derive(Debug, Clone)]
pub struct TypeConstraint {
    /// Left side of constraint
    pub left: Type,
    /// Right side of constraint
    pub right: Type,
    /// Constraint location
    pub location: crate::error::Location,
}

/// Tracks the enclosing return type that can consume propagated errors
#[derive(Clone)]
enum TryContext {
    /// Function/lambda returning Option<T>
    Option(crate::ast::types::OptionType),
    /// Function/lambda returning Result<T, E>
    Result(crate::ast::types::ResultType),
}

impl TryContext {
    fn from_type(ty: &Type) -> Option<Self> {
        match ty {
            Type::Option(option_type) => Some(TryContext::Option(option_type.clone())),
            Type::Result(result_type) => Some(TryContext::Result(result_type.clone())),
            _ => None,
        }
    }
}

impl TypeInference {
    /// Create a new type inference engine
    pub fn new() -> Self {
        Self {
            environment: TypeEnvironment::new(),
            constraints: Vec::new(),
            inferred_types: HashMap::new(),
             try_context_stack: Vec::new(),
        }
    }

    /// Add a type constraint
    fn add_constraint(
        &mut self,
        left: &Type,
        right: &Type,
        location: crate::error::Location,
    ) -> Result<()> {
        self.constraints.push(TypeConstraint {
            left: left.clone(),
            right: right.clone(),
            location,
        });
        Ok(())
    }

    /// Resolve a type, converting generic types to their actual definitions if they exist
    fn resolve_type(&self, type_: &Type) -> Result<Type> {
        match type_ {
            Type::Generic(name) => {
                if let Some(defined_type) = self.environment.lookup_type(name) {
                    self.resolve_type(defined_type)
                } else {
                    Ok(type_.clone())
                }
            }
            Type::Struct(struct_type) => {
                let mut resolved_type_args = Vec::new();
                for type_arg in &struct_type.type_args {
                    resolved_type_args.push(self.resolve_type(type_arg)?);
                }
                Ok(Type::Struct(StructType {
                    name: struct_type.name.clone(),
                    type_args: resolved_type_args,
                    location: struct_type.location,
                }))
            }
            Type::Result(result_type) => {
                let ok_type = self.resolve_type(&result_type.ok_type)?;
                let err_type = self.resolve_type(&result_type.err_type)?;
                Ok(Type::Result(crate::ast::types::ResultType {
                    ok_type: Box::new(ok_type),
                    err_type: Box::new(err_type),
                    location: result_type.location,
                }))
            }
            _ => Ok(type_.clone()),
        }
    }

    /// Check if the type inference engine is empty
    pub fn is_empty(&self) -> bool {
        self.constraints.is_empty() && self.inferred_types.is_empty()
    }

    /// Get the type environment
    pub fn environment(&self) -> &TypeEnvironment {
        &self.environment
    }

    /// Get the constraints
    pub fn constraints(&self) -> &Vec<TypeConstraint> {
        &self.constraints
    }

    /// Get the inferred types
    pub fn inferred_types(&self) -> &HashMap<String, Type> {
        &self.inferred_types
    }

    /// Push a new try context derived from a return type
    fn push_try_context_from_type(&mut self, ty: Option<&Type>) {
        let ctx = ty.and_then(TryContext::from_type);
        self.try_context_stack.push(ctx);
    }

    /// Restore the try-context stack to a previous depth
    fn truncate_try_context(&mut self, depth: usize) {
        self.try_context_stack.truncate(depth);
    }

    /// Get the current active try context, if any
    fn current_try_context(&self) -> Option<&TryContext> {
        self.try_context_stack
            .last()
            .and_then(|ctx| ctx.as_ref())
    }
}

impl Default for TypeInference {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_type_inference_creation() {
        let inference = TypeInference::new();
        assert!(inference.is_empty());
    }

    #[test]
    fn test_literal_type_inference() {
        let mut inference = TypeInference::new();
        let mut env = TypeEnvironment::new();

        let int_lit = Expression::Literal(Literal::Integer(42));
        let int_type = inference.infer_expression(&int_lit, &mut env).unwrap();
        assert_eq!(int_type, Type::Basic(crate::ast::types::BasicType::Int));

        let str_lit = Expression::Literal(Literal::String("hello".to_string()));
        let str_type = inference.infer_expression(&str_lit, &mut env).unwrap();
        assert_eq!(str_type, Type::Basic(crate::ast::types::BasicType::String));
    }

    #[test]
    fn test_variable_type_inference() {
        let mut inference = TypeInference::new();
        let mut env = TypeEnvironment::new();

        let int_type = Type::Basic(crate::ast::types::BasicType::Int);
        env.bind_variable("x".to_string(), int_type.clone());

        let var_expr = Expression::Variable("x".to_string());
        let inferred_type = inference.infer_expression(&var_expr, &mut env).unwrap();
        assert_eq!(inferred_type, int_type);
    }

    #[test]
    fn test_binary_expression_type_inference() {
        let mut inference = TypeInference::new();
        let mut env = TypeEnvironment::new();

        let left = Expression::Literal(Literal::Integer(10));
        let right = Expression::Literal(Literal::Integer(20));
        let binary = Expression::Binary(BinaryExpr {
            left: Box::new(left),
            operator: crate::ast::expr::BinaryOp::Add,
            right: Box::new(right),
            location: crate::error::Location::new(1, 1, 0),
        });

        let inferred_type = inference.infer_expression(&binary, &mut env).unwrap();
        assert_eq!(
            inferred_type,
            Type::Basic(crate::ast::types::BasicType::Int)
        );
    }
}
